#include "Agent.h"

#include <unistd.h>

#include <iostream>

#include "dss_common.h"
using namespace std;
struct AgentContext
{
    RpcExecution *cli2agent_ex;
    uint64_t cli2agent_sign;
    uint32_t cli2agent_sync; 
    uint32_t cli2agent_peer_id;

    RpcExecution *agent2server_ex;

    Agent *agent;
};

static void 
_forward_message_cb(void *arg1, void *arg2, void* arg3)
{
    RdmaRpc *rpc = static_cast<RdmaRpc*>(arg1);
    RpcExecution *ex = static_cast<RpcExecution*>(arg2);
    AgentContext *ctx = static_cast<AgentContext*>(arg3);
    Agent *agent = ctx->agent;

    //at ring buf 
    RpcMsg *rpc_resp = ex->resp_msg;
    RpcExecution *client2agent_ex = ctx->cli2agent_ex;


    RpcMsg *client2agent_req = client2agent_ex->req_msg;
    RpcMsg *client2agent_resp = client2agent_ex->resp_msg;

    /* recover request ctx*/
    client2agent_req->sign = ctx->cli2agent_sign;
    client2agent_req->sync = ctx->cli2agent_sync;
    client2agent_req->peer_id = ctx->cli2agent_peer_id;

    rpc->FreeMsg(client2agent_resp);
    client2agent_ex->resp_msg = rpc_resp;
    client2agent_ex->wait_free_sess_id = rpc_resp->peer_id;
    client2agent_ex->wait_free_rpc = rpc;
    //memcpy(client2agent_ex->resp_msg, rpc_resp, sizeof(RpcMsg) + sizeof(dss_msg) + 64);
    ex->resp_msg = nullptr;
    
    ctx->agent->rpc_server_.RpcServiceComplete(client2agent_ex);
    rpc->FreePiece(ctx);
    

}
static void
_forward_message(void *arg1, void *arg2, void* arg3)
{
    RdmaRpc *rpc = static_cast<RdmaRpc*>(arg1);
    RpcExecution *ex = static_cast<RpcExecution*>(arg2);
    Agent *agent = static_cast<Agent*>(arg3);

    AgentContext *ctx = static_cast<AgentContext*>(rpc->AllocPiece());   
    RpcMsg *rpc_resp = ex->resp_msg;
    RpcMsg *rpc_req = ex->req_msg;

    /*store request ctx*/
    ctx->cli2agent_ex = ex;
    ctx->cli2agent_peer_id = rpc_req->peer_id;
    ctx->cli2agent_sync = rpc_req->sync;
    ctx->cli2agent_sign = rpc_req->sign;
    ctx->agent2server_ex = static_cast<RpcExecution*>(rpc->AllocPiece());

    ctx->agent = agent;

    /*switch request ctx*/
    rpc_req->sign = reinterpret_cast<uint64_t>(ctx->agent2server_ex);
    ctx->agent2server_ex->req_msg = rpc_req;
    ctx->agent2server_ex->resp_msg = NULL;
    agent->rpc_cli_.RpcCliWriteReq(agent->cs_sess_id_, rpc_req, _forward_message_cb, ctx);

}

int 
main(int argc, char** argv) {
	int rc = 0;
	char ch;
    char *local_uri;
    char *server_uri;
	while ((ch = getopt(argc, argv, "l:r:h")) != -1) {
		switch (ch) {
		case 'l':
			local_uri = strdup(optarg);
			break;
		case 'r':
			server_uri = strdup(optarg);
			break;
		case 'h':
			cout << "\n-l Agent local uri \n-r Server uri\n-h Help\n";
			return rc;
		default:
			cout << "NOT recognized option: " << ch << endl;
			return rc;
		}
	}
    uint32_t mem_size = 2 << 30;
    void* mem = new char[mem_size];
    
    MemoryManager mm(mem, mem_size, 256 << 10);
    mm.Init();

    RdmaRpc cli(mm);
    cli.Init(RPC_RDMA | RPC_SHARED_RECV_QUEUE);
    

    RdmaRpc server(mm);
    server.Init(RPC_SERVER | RPC_RDMA | RPC_SHARED_RECV_QUEUE);



    Agent agent(mm, cli, server);

    agent.cs_sess_id_ = cli.ConnectSess(server_uri, nullptr, nullptr);

    server.RpcServiceRegister(S_CHUNK_IO_SERVICE, _forward_message);
    server.RpcServiceProviderRegister(S_CHUNK_IO_SERVICE, &agent);
    server.BindSess(local_uri);
    uint64_t last_active = 0;
    while(1)
    {
        server.RpcServerHandleRx();
        server.RpcServerHandleTx();
        cli.RpcCliHandleRx();
        cli.RpcCliHandleTx();
        uint64_t cur = Cycles::rdtsc();
        if(Cycles::to_seconds(cur - last_active, 0) > 1.0)
        {
            server.AcceptSess();
            last_active = cur;
        }
    }
    
}